from collections import namedtuple

import pandas as pd
import scipy.stats as stats 
from statsmodels.stats.multitest import multipletests


def preprocess(
    owwwanalyzer_output_excel_filepath: str, 
    sample_group_info_filepath: str
) -> dict[str, pd.DataFrame]:

    colnames = namedtuple(
        "ColNames", ["sample_id", "group", "assay", "npx"]
    )("SampleID", "Group", "Assay", "NPX")
    
    full_result = pd.read_excel(owwwanalyzer_output_excel_filepath, "原始数据")
    all_assays = full_result[colnames.assay].unique()
    
    if sample_group_info_filepath.endswith(".csv"):
        group_info = pd.read_csv(sample_group_info_filepath, header=0)
    elif sample_group_info_filepath.endswith(".xlsx"):
        group_info = pd.read_excel(sample_group_info_filepath)
    else:
        raise ValueError("Only .csv or .xlsx file formats are supported")
    
    if any(
        col not in group_info.columns 
        for col in (colnames.sample_id, colnames.group)
    ):

        raise ValueError(
            "The sample group information table must "
            "contain the columns 'SampleID' and 'Group'"
        )

    grouped_npx_dfs = dict()
    for label, df in group_info.groupby(colnames.group):
        sample_ids = df[colnames.sample_id].tolist()
        npx_df = full_result.loc[
            full_result[colnames.sample_id].isin(sample_ids), 
            [colnames.sample_id, colnames.assay, colnames.npx]
        ].pivot(
            index=colnames.assay, 
            columns=colnames.sample_id, 
            values=colnames.npx
        ).reindex(all_assays, axis="index")
        grouped_npx_dfs[label] = npx_df

    return grouped_npx_dfs


def calculate_deps_with_ranksum(
    control_npx_df: pd.DataFrame,
    treat_npx_df: pd.DataFrame, 
    fdr_threshold: float=0.05
) -> pd.DataFrame:

    assert (control_npx_df.index == treat_npx_df.index).all()
    
    de_results = pd.DataFrame(index=control_npx_df.index)
    de_results["log2FC"] = (
        treat_npx_df.median(axis=1) - control_npx_df.median(axis=1)
    )

    test_result = stats.mannwhitneyu(
        treat_npx_df, control_npx_df, nan_policy="omit", axis=1
    )
    de_results["pvalue"] = test_result.pvalue
    
    # 多重检验校正
    de_results["significant"], de_results["fdr"], *_ = multipletests(
        de_results["pvalue"], alpha=fdr_threshold, method="fdr_bh"
    )
    
    de_results["regulation"] = "NoDiff"
    up_mask = (de_results["log2FC"] > 0) & (de_results["significant"])
    down_mask = (de_results["log2FC"] < 0) & (de_results["significant"])
    de_results.loc[up_mask, "regulation"] = "Sig_Up"
    de_results.loc[down_mask, "regulation"] = "Sig_Down"
    
    return de_results
